#define _CRT_SECURE_NO_WARNINGS
#include <Shobjidl.h>
#include <ObjIdl.h>
#define USE_SYS_LIBC
//#define USE_SYS_MSVCRT
#include "../ntlib/ntlib.hpp"
#include "lsap.h"
extern shared_struct* _shared;
extern WCHAR moduleName[MAX_PATH];

#pragma comment(lib, "ole32.lib")
struct s_ptr
{
    s_ptr()
    {
        ntlib::debug_printw(L"s_ptr ctor");
        _shared->mode  = 322;
    }
    ~s_ptr()
    {
        ntlib::debug_printw(L"s_ptr dctor");
        _shared->mode = 32422222;
    }
};

typedef struct IContextMenuVtbl
{
    BEGIN_INTERFACE

        HRESULT ( STDMETHODCALLTYPE *QueryInterface )( 
        IContextMenu * This,
        /* [in] */ REFIID riid,
        /* [annotation][iid_is][out] */ 
        _COM_Outptr_  void **ppvObject);

        ULONG ( STDMETHODCALLTYPE *AddRef )( 
            IContextMenu * This);

        ULONG ( STDMETHODCALLTYPE *Release )( 
            IContextMenu * This);

        HRESULT ( STDMETHODCALLTYPE *QueryContextMenu )( 
            IContextMenu * This,
            /* [annotation][in] */ 
            _In_  HMENU hmenu,
            /* [annotation][in] */ 
            _In_  UINT indexMenu,
            /* [annotation][in] */ 
            _In_  UINT idCmdFirst,
            /* [annotation][in] */ 
            _In_  UINT idCmdLast,
            /* [annotation][in] */ 
            _In_  UINT uFlags);

        HRESULT ( STDMETHODCALLTYPE *InvokeCommand )( 
            IContextMenu * This,
            /* [annotation][in] */ 
            _In_  CMINVOKECOMMANDINFO *pici);

        HRESULT ( STDMETHODCALLTYPE *GetCommandString )( 
            IContextMenu * This,
            /* [annotation][in] */ 
            _In_  UINT_PTR idCmd,
            /* [annotation][in] */ 
            _In_  UINT uType,
            /* [annotation][in] */ 
            _Reserved_  UINT *pReserved,
            /* [annotation][out] */ 
            _Out_writes_bytes_((uType & GCS_UNICODE) ? (cchMax * sizeof(wchar_t)) : cchMax) _When_(!(uType & (GCS_VALIDATEA | GCS_VALIDATEW)), _Null_terminated_)  CHAR *pszName,
            /* [annotation][in] */ 
            _In_  UINT cchMax);

    END_INTERFACE
} IContextMenuVtbl;

struct IContextMenu_C
{
    CONST_VTBL struct IContextMenuVtbl *lpVtbl;
};

// {00021401-0000-0000-C000-000000000046}
static const GUID shellLinkCLSID = 
{ 0x00021401, 0x0, 0x0, { 0xC0, 0, 0, 0, 0, 0, 0, 0x46 } };

LPVOID shLinkObj()
{
    bool CoInit = false;
    LPVOID invoke = 0;
_start:
    IUnknown* obj = 0;
    auto hr = CoCreateInstance(shellLinkCLSID,0,CLSCTX_INPROC_SERVER,IID_IUnknown,(LPVOID*)&obj);
    if (hr == CO_E_NOTINITIALIZED)
    {
        CoInitialize(0);
        CoInit = true;
        goto _start;
    }
    obj->AddRef();
    IContextMenu_C* menu;
    hr = obj->QueryInterface(IID_IContextMenu,(LPVOID*)&menu);
    invoke = menu->lpVtbl->InvokeCommand;
    obj->Release();
    if (CoInit)
        CoUninitialize();
    return invoke;
}

int launch(LPCWSTR path,LPCWSTR params, LPCWSTR descr, LPCWSTR workdir)
{
    ntlib::debug_printw(L"path: %s",path);
    ntlib::debug_printw(L"param: %s",params);
    ntlib::debug_printw(L"desc: %s",descr);
    ntlib::debug_printw(L"workdir: %s",workdir);
    if (lstrcmpiW(descr,L"limit") == 0 || descr[0] == L'!')
    {
        WCHAR tmp [MAX_PATH * 2 + 8] = L"\"";
        wcscat(tmp,path);
        wcscat(tmp,L"\" ");
        wcscat(tmp,params);

        if (descr[0] == L'!')
        {
            WCHAR p[MAX_PATH];
            wcscpy(p,descr+1);
            auto pass =wcsstr(p,L" ");
            if (!pass) return 0;
            *pass++ = 0;
            PROCESS_INFORMATION pinfo = {};
            STARTUPINFO sinfo = {sizeof(STARTUPINFO)};
            auto ret = CreateProcessWithLogonW(p,0,pass,LOGON_WITH_PROFILE,path,tmp,
                0,0,workdir,&sinfo,&pinfo);
            if (ret)
            {
                CloseHandle(pinfo.hProcess);
                CloseHandle(pinfo.hThread);
                return 1;
            }
            return 0;
        }
        else
        {
            static auto lsass_id = ntlib::find_process_id(L"lsass.exe");
            _shared->do_resident = 0;
            if (ntlib::inject_me(lsass_id,moduleName))
            {
                unsigned long WINAPI limit_token(PVOID p);
                HANDLE WINAPI restrict_token(PVOID);
                PROCESS_INFORMATION pinfo = {};
                STARTUPINFO sinfo = {sizeof(STARTUPINFO)};
                auto t = restrict_token(0);
                auto ret = CreateProcessAsUserW(t,path,tmp,
                    0,0,0,CREATE_SUSPENDED,0,workdir,&sinfo,&pinfo);
                CloseHandle(t);
                if (ret)
                {
                    ntlib::call_remote(lsass_id,limit_token,(void*)pinfo.dwProcessId);
                    ResumeThread(pinfo.hThread);
                    CloseHandle(pinfo.hProcess);
                    CloseHandle(pinfo.hThread);
                }
                ntlib::call_remote(lsass_id,[](PVOID) -> unsigned long{
                    extern HMODULE Module;
                    FreeLibraryAndExitThread(Module,0);
                    return 0;
                },0);
                if (ret)
                    return 1;
            }
        }
    }
    return 0;
}
HRESULT ( STDMETHODCALLTYPE *real_invoke )( 
    IContextMenu * This,
    /* [annotation][in] */ 
    _In_  CMINVOKECOMMANDINFO *pici);

HRESULT ( STDMETHODCALLTYPE InvokeCommand )( 
    IContextMenu * This,
    /* [annotation][in] */ 
    _In_  CMINVOKECOMMANDINFO *pici)
{
    char tmp[32] = {};
    const char * target = pici->lpVerb;
    if (IS_INTRESOURCE(pici->lpVerb))
    {
        This->GetCommandString((UINT_PTR)pici->lpVerb,GCS_VERBA,0,tmp,32);
        target = tmp;
    }
    if (lstrcmpiA(target,"open") == 0)
    {
        IShellLinkW* sL = 0;
        //IPersistFile* sF = 0;
        This->QueryInterface(IID_IShellLinkW,(LPVOID*)&sL);
        //This->QueryInterface(IID_IPersistFile,(LPVOID*)&sF);
        if (sL)
        {
            sL->AddRef();
            //sF->AddRef();
            WCHAR tmp[MAX_PATH] = {}, params[MAX_PATH] = {},des[MAX_PATH] = {},workdir[MAX_PATH]={};
            //LPOLESTR file = 0;
            sL->GetPath(tmp,MAX_PATH,0,0);
            sL->GetDescription(des,MAX_PATH);
            sL->GetArguments(params,MAX_PATH);
            sL->GetWorkingDirectory(workdir,MAX_PATH);
            //sF->GetCurFile(&file);
            sL->Release();
            //sF->Release();
            int r = launch(tmp,params,des,workdir);
            if (r)
            {
                return S_OK;
            }
        }
    }
    return real_invoke(This,pici);
}
static ntlib::hook_info shell_hook_info = {0};
DWORD WINAPI hook_shell_remote(LPVOID)
{
    auto invoke = shLinkObj();
    if (ntlib::hook_proc(invoke,InvokeCommand,&real_invoke,13,&shell_hook_info))
    {
        ntlib::debug_printw(L"hook shell %x",invoke);
    }
    return 1;
}

void hook_shell(WCHAR* pname)
{
    auto pid = ntlib::find_process_id(pname);
    if (pid)
    {
        ntlib::debug_printw(L"shell hooking");
        _shared->do_resident = 0;
        if (ntlib::inject_me(pid,moduleName))
        {
            ntlib::call_remote(pid,hook_shell_remote,0);
        }
    }
}
void free_shell(WCHAR* pname)
{
    auto pid = ntlib::find_process_id(pname);
    if (pid)
    {
        ntlib::call_remote(pid,[](PVOID) -> unsigned long{
            if (ntlib::unhook_proc(&shell_hook_info))
            {
                ntlib::debug_printw(L"unhook shell");
            }
            extern HMODULE Module;
            FreeLibraryAndExitThread(Module,0);
            return 0;
        },0);
    }
}